#!/usr/bin/env python3
import argparse
from flask import Flask, render_template, request     # 🟩【修改】新增 request
from flask_socketio import SocketIO
import pty
import os
import subprocess
import select
import termios
import struct
import fcntl
import shlex
import logging
import sys
import re
import threading           # 🟩【新增】支持锁
import signal              # 🟩【新增】回收僵尸进程
import time                # 🟩【新增】守护线程周期

logging.getLogger("werkzeug").setLevel(logging.ERROR)
__version__ = "0.5.2.0"

# ------------------- Flask 初始化 -------------------
app = Flask(__name__, template_folder=".", static_folder=".", static_url_path="")
app.config["SECRET_KEY"] = "secret!"
app.config["cmd"] = None

socketio = SocketIO(app)

# ------------------- 多客户端管理 -------------------
clients = {}                   # { sid: {"fd": fd, "child_pid": pid} }
clients_lock = threading.Lock()  # 🟩【新增】防止多线程访问冲突

# ------------------- 工具函数 -------------------
# def set_winsize(fd, row, col, xpix=0, ypix=0):
#     winsize = struct.pack("HHHH", row, col, xpix, ypix)
#     fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)

ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')

# ------------------- 信号处理：自动回收僵尸进程 -------------------
def reap_children(signum, frame):   # 🟩【新增】
    """SIGCHLD 处理，回收子进程避免僵尸进程"""
    try:
        while True:
            pid, _ = os.waitpid(-1, os.WNOHANG)
            if pid == 0:
                break
    except ChildProcessError:
        pass

signal.signal(signal.SIGCHLD, reap_children)  # 🟩【新增】

# ------------------- 后台读取任务 -------------------
def read_and_forward_pty_output(sid):
    """后台线程：持续读取子进程输出并发送前端"""
    max_read_bytes = 1024 * 20
    while True:
        socketio.sleep(0.02)
        with clients_lock:
            if sid not in clients:
                break
            fd = clients[sid]["fd"]
        data_ready, _, _ = select.select([fd], [], [], 0)
        if data_ready:
            try:
                data = os.read(fd, max_read_bytes).decode('utf-8', errors="ignore")
                output = ansi_escape.sub('', data)
                socketio.emit("pty-output", {"output": output}, namespace="/pty", to=sid)
            except OSError:
                break

# ------------------- Flask 路由 -------------------
@app.route("/")
def index():
    return render_template("index.html")

# ------------------- SocketIO事件 -------------------
@socketio.on("pty-input", namespace="/pty")
def pty_input(data):
    """前端输入写入子进程"""
    sid = request.sid
    with clients_lock:
        if sid not in clients:
            return
        fd = clients[sid]["fd"]
    os.write(fd, data["input"].encode())

# @socketio.on("resize", namespace="/pty")
# def resize(data):
#     """调整终端大小"""
#     sid = request.sid
#     with clients_lock:
#         if sid not in clients:
#             return
#         fd = clients[sid]["fd"]
#     set_winsize(fd, data["rows"], data["cols"])

# 🟩【修改】connect 接受可选参数避免 TypeError
@socketio.on("connect", namespace="/pty")
def connect(*args, **kwargs):
    """新客户端连接"""
    sid = request.sid
    logging.info(f"new client connected: {sid}")

    # ---------------- 子进程逻辑 ----------------
    child_pid, fd = pty.fork()
    if child_pid == 0:
        # ---------------- 子进程 ----------------
        # 切换到部署用户 HOME
        home = os.path.expanduser("~")
        os.chdir(home)

        # 清理虚拟环境：删除 VIRTUAL_ENV
        os.environ.pop("VIRTUAL_ENV", None)

        # 清理 PATH 中的虚拟环境路径
        if 'PATH' in os.environ:
            venv_bin = os.path.dirname(sys.executable)  # /path/to/.venv/bin
            paths = os.environ['PATH'].split(":")
            paths = [p for p in paths if p != venv_bin]
            os.environ['PATH'] = ":".join(paths)

        # os.setsid()  # 🟥【删除】在 Eventlet 下容易 PermissionError
        subprocess.run(["/bin/bash", "--norc", "--noprofile"], env={**os.environ, "PS1": "$ "})
        # 启动干净的 bash，不继承 .bashrc / profile
        # os.execvp("bash", ["bash", "--noprofile", "--norc"])
        os._exit(0)  # 🟩【新增】子进程正常退出
    else:
        # ---------------- 父进程逻辑 ----------------
        with clients_lock:
            clients[sid] = {"fd": fd, "child_pid": child_pid}
        # set_winsize(fd, 50, 50)
        socketio.start_background_task(target=read_and_forward_pty_output, sid=sid)
        logging.info(f"started bash pid={child_pid} for sid={sid}")

@socketio.on("disconnect", namespace="/pty")
def disconnect():
    """客户端断开，清理资源"""
    sid = request.sid
    with clients_lock:
        client = clients.pop(sid, None)
    if client:
        try:
            os.close(client["fd"])
        except Exception:
            pass
        try:
            os.kill(client["child_pid"], 9)
        except ProcessLookupError:
            pass
    logging.info(f"client {sid} disconnected, cleaned up")

# ------------------- 守护线程：清理死进程记录 -------------------
def cleanup_task():   # 🟩【新增】
    """后台任务检测 /proc 子进程是否存在"""
    while True:
        time.sleep(30)
        with clients_lock:
            dead_sids = []
            for sid, info in list(clients.items()):
                if not os.path.exists(f"/proc/{info['child_pid']}"):
                    dead_sids.append(sid)
            for sid in dead_sids:
                logging.info(f"cleanup: removing stale sid={sid}")
                del clients[sid]

# 🟩【新增】启动后台清理任务
threading.Thread(target=cleanup_task, daemon=True).start()

# ------------------- 主入口 -------------------
def main():
    parser = argparse.ArgumentParser(description="Enhanced pyxterm.js Server")
    parser.add_argument("-p", "--port", default=5000, type=int)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()

    green, end = "\033[92m", "\033[0m"
    log_format = green + "pyxtermjs > " + end + "%(levelname)s (%(funcName)s:%(lineno)s) %(message)s"
    logging.basicConfig(format=log_format, stream=sys.stdout, level=logging.DEBUG if args.debug else logging.INFO)

    logging.info(f"serving on http://{args.host}:{args.port}")
    socketio.run(app, debug=args.debug, port=args.port, host=args.host)

if __name__ == "__main__":
    main()
